#include <BWTA.h>
#include "HarassmentManager.h"
#include <UnitGroupManager.h>
using namespace BWAPI;


HarassmentManager::HarassmentManager(Arbitrator::Arbitrator<BWAPI::Unit*,double> *arbitrator)
{
	this->arbitrator = arbitrator;
	desiredHarasserCount = 1;
	//TODO:  Would work better if it could reserve the resources it needed. (Through the BuildOrderManager)
	//			Perhaps try std::vector<BuildData> builds;
	//			Where BuildData stores the location of the building, the building type, and a flag for if minerals have been reserved.
	//			If you can build in the location (there is space), then reserve the resources (whether you have them or not) and (if you have the resources) build.
	//				Only do this for the first building in the list which can be built (so it doesn't tie up ALL of the resources).
	//			Keep only the next thing ot be built reserved.
}

void HarassmentManager::onOffer(std::set<BWAPI::Unit*> units)
{
	std::set<BWAPI::Unit*>::iterator u2;

	for(std::set<BWAPI::Unit*>::iterator u = units.begin(); u != units.end(); u=u2)
	{
		u2=u;
		u2++;
		//ignore if its already a harasser
		if (harassers.find(*u) != harassers.end())
		{
			arbitrator->accept(this, *u);
			units.erase(u);
			continue;
		}
		if ((*u)->getType().isWorker() && needMoreHarassers())
		{
			arbitrator->accept(this, *u);
			addHarasser(*u);
			units.erase(u);
		}
	}

	//decline remaining units
	for(std::set<BWAPI::Unit*>::iterator u = units.begin(); u != units.end(); u++)
	{
		arbitrator->decline(this, *u, 0);
	}
}


void HarassmentManager::update()
{
	if (needMoreHarassers())
	{
		requestHarasser(/* bid = */ 40);
	}
	else
	{
		while (harassers.size()>desiredHarasserCount)
		{
			arbitrator->setBid(this, harassers.begin()->first,0);
			harassers.erase(harassers.begin());
		}
	}
	updateHarasserAssignments();
}

std::string HarassmentManager::getName() const
{
	return "Harassment Manager";
}

std::string HarassmentManager::getShortName() const
{
	return "Harass";
}

void HarassmentManager::setHarasserCount(int count)
{
	this->desiredHarasserCount=count;
}

bool HarassmentManager::isHarassing() const
{
	return harassers.size() >= 1;
}

bool HarassmentManager::needMoreHarassers() const
{
	return harassers.size() < desiredHarasserCount;
}

void HarassmentManager::requestHarasser(double bid)
{
	// Bid on all completed workers.
	std::set<BWAPI::Unit*> usefulUnits=SelectAll()(isWorker)(isCompleted).not(isCarryingMinerals,isCarryingGas,isGatheringGas);
	arbitrator->setBid(this,usefulUnits,bid);
}

void HarassmentManager::addHarasser(BWAPI::Unit* u)
{
	HarasserData temp;
	harassers.insert(std::make_pair(u,temp));
}

void HarassmentManager::onRevoke(BWAPI::Unit *unit, double bid)
{
	onRemoveUnit(unit);
}

void HarassmentManager::onRemoveUnit(BWAPI::Unit* unit)
{
	if (harassers.find(unit) != harassers.end())
	{
		harassers.erase(unit);
	}
}

void HarassmentManager::updateHarasserAssignments()
{
	Position pos;	

	bool abort = false;
	/*bool have_photon = false;
	bool first_pylon_under_construction = false;	

	std::set<Unit*> enemyUnits;
	for(std::set<Unit*>::iterator m = Broodwar->getAllUnits().begin();m != Broodwar->getAllUnits().end(); ++m)
	{
	if ((*m)->getPlayer()->isEnemy(Broodwar->self()))
	{
	enemyUnits.insert((*m));
	}
	}

	for(std::set<Unit*>::const_iterator i = Broodwar->self()->getUnits().begin(); i != Broodwar->self()->getUnits().end(); i++){
	if ((*i)->getID() == 162){
	have_photon = true;
	}
	}

	for(std::set<Unit*>::const_iterator i = Broodwar->self()->getUnits().begin(); i != Broodwar->self()->getUnits().end(); i++){
	if (((*i)->getID() == 156 || (*i)->getID() == 166) && !enemyUnits.empty() && !have_photon){ //pylon tai forge
	for(std::set<Unit*>::const_iterator j = enemyUnits.begin(); j != enemyUnits.end(); j++){
	if ((*j)->getTarget() == (*i)){
	abort = true;
	}
	}
	}
	if ((*i)->getID() == 156){ //pylon
	if ((*i)->isBeingConstructed()){
	first_pylon_under_construction = true;
	}
	}
	}*/

	//if (Broodwar->getFrameCount() > 5200){
	//	abort = true;
	//}

	if (abort){
		setHarasserCount(0);
	}

	if(Broodwar->self()->getStartLocation() == Position(63*32, 117*32))
	{
		std::map<BWAPI::Unit*, HarasserData>::iterator u;
		for(u = harassers.begin(); u != harassers.end(); u++)
		{
			pos = Position(32*56, 32*0);	 
			if ((*u).second.mode == HarasserData::Idle)
			{
				(*u).first->rightClick(pos);
				(*u).second.mode = HarasserData::Harassing;
			}
			if (BWAPI::Broodwar->canBuildHere((*u).first, pos, UnitTypes::Protoss_Pylon) && Broodwar->self()->minerals() >= 100)
				(*u).first->build(pos, UnitTypes::Protoss_Pylon);

			//else if (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x()-96, 0), UnitTypes::Protoss_Forge) && Broodwar->self()->minerals() >= 150)
			//	(*u).first->build(Position(pos.x()-96, 0), UnitTypes::Protoss_Forge);

			else if (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() - 64, 32), UnitTypes::Protoss_Photon_Cannon) && Broodwar->self()->minerals() >= 150)
				(*u).first->build(Position(pos.x() - 64, 32), UnitTypes::Protoss_Photon_Cannon);

			else if (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() - 64, 32), UnitTypes::Protoss_Photon_Cannon) && Broodwar->self()->minerals() >= 150)
				(*u).first->build(Position(pos.x() +  0, 64), UnitTypes::Protoss_Photon_Cannon);

			else if (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 64, 32), UnitTypes::Protoss_Photon_Cannon) && Broodwar->self()->minerals() >= 150)
				(*u).first->build(Position(pos.x() + 64, 32), UnitTypes::Protoss_Photon_Cannon);

			else if  (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 128, 32), UnitTypes::Protoss_Photon_Cannon) && Broodwar->self()->minerals() >= 150)
				(*u).first->build(Position(pos.x() + 128, 32), UnitTypes::Protoss_Photon_Cannon);

			else if  (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 192, 32), UnitTypes::Protoss_Photon_Cannon ) && Broodwar->self()->minerals() >= 150)
				(*u).first->build(Position(pos.x() + 192, 32), UnitTypes::Protoss_Photon_Cannon);

			else if (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 256, 0), UnitTypes::Protoss_Pylon) && Broodwar->self()->minerals() >= 200)
				(*u).first->build(Position(pos.x() + 256, 0), UnitTypes::Protoss_Pylon);

			else if  (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 256, 64), UnitTypes::Protoss_Photon_Cannon) && Broodwar->self()->minerals() >= 150)
				(*u).first->build(Position(pos.x() + 256, 64), UnitTypes::Protoss_Photon_Cannon);

			else if  (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 320, 64), UnitTypes::Protoss_Photon_Cannon) && Broodwar->self()->minerals() >= 150){
				(*u).first->build(Position(pos.x() + 320, 64), UnitTypes::Protoss_Photon_Cannon);
				onRemoveUnit((*u).first);
				setHarasserCount(0);
			}
			else  
				(*u).first->rightClick(Position(pos.x() + 64, 0));

		}
	}

	else if (Broodwar->self()->getStartLocation() == Position(70*32, 8*32))
	{

		std::map<BWAPI::Unit*, HarasserData>::iterator u;

		for(u = harassers.begin(); u != harassers.end(); u++)
		{

			pos = Position(46*32, 125*32);	 

			if ((*u).second.mode == HarasserData::Idle)
			{
				(*u).first->rightClick(pos);
				(*u).second.mode = HarasserData::Harassing;
			}

			if (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 64, pos.y()), UnitTypes::Protoss_Pylon) && Broodwar->self()->minerals() >= 100)
				(*u).first->build(Position(pos.x() + 64, pos.y()), UnitTypes::Protoss_Pylon);

			//else if (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 64, pos.y() - 64), UnitTypes::Protoss_Forge) && Broodwar->self()->minerals() >= 150)
			//	(*u).first->build(Position(pos.x() + 64, pos.y() - 64), UnitTypes::Protoss_Forge);

			else if (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 64 + 64, pos.y()), UnitTypes::Protoss_Photon_Cannon) && Broodwar->self()->minerals() >= 150)
				(*u).first->build(Position(pos.x() +  0, pos.y() - 64), UnitTypes::Protoss_Photon_Cannon);

			else if (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 64 + 64, pos.y()), UnitTypes::Protoss_Photon_Cannon) && Broodwar->self()->minerals() >= 150)
				(*u).first->build(Position(pos.x() + 64 + 64, pos.y()), UnitTypes::Protoss_Photon_Cannon);

			else if  (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 64 + 64 + 64, pos.y()), UnitTypes::Protoss_Photon_Cannon) && Broodwar->self()->minerals() >= 150)
				(*u).first->build(Position(pos.x() + 64 + 64 + 64, pos.y()), UnitTypes::Protoss_Photon_Cannon);

			else if  (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 64 + 64 + 64 + 64, pos.y()), UnitTypes::Protoss_Photon_Cannon ) && Broodwar->self()->minerals() >= 150)
				(*u).first->build(Position(pos.x() + 64 + 64 + 64 + 64, pos.y()), UnitTypes::Protoss_Photon_Cannon);


			else if (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 64 + 64 + 64 + 64 + 64, pos.y()), UnitTypes::Protoss_Pylon) && Broodwar->self()->minerals() >= 200)
				(*u).first->build(Position(pos.x() + 64 + 64 + 64 + 64 +64, pos.y()), UnitTypes::Protoss_Pylon);

			else if  (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 64 + 64 + 64 + 64 + 64, pos.y() - 64), UnitTypes::Protoss_Photon_Cannon) && Broodwar->self()->minerals() >= 150)
				(*u).first->build(Position(pos.x() + 64 + 64 + 64 + 64 + 64, pos.y() - 64), UnitTypes::Protoss_Photon_Cannon);

			else if  (BWAPI::Broodwar->canBuildHere((*u).first, Position(pos.x() + 64 + 64 + 64 + 64 + 64 + 64, pos.y() - 64), UnitTypes::Protoss_Photon_Cannon) && Broodwar->self()->minerals() >= 150){
				(*u).first->build(Position(pos.x() + 64 + 64 + 64 + 64 + 64 + 64, pos.y() - 64), UnitTypes::Protoss_Photon_Cannon);
				onRemoveUnit((*u).first);
				setHarasserCount(0);
			}
			else  
				(*u).first->rightClick(Position(pos));
		}
	} else {
		setHarasserCount(0);
	}
}